#include "receiver_harness.hh"
#include "util.hh"
#include "wrapping_integers.hh"

#include <cstdint>
#include <cstdlib>
#include <exception>
#include <iostream>
#include <optional>
#include <random>
#include <stdexcept>
#include <string>

using namespace std;

int main() {
  try {
    auto rd = get_random_generator();

    {
      TCPReceiverTestHarness test{4000};
      test.execute(SegmentArrives{}.with_syn().with_seqno(0).with_result(SegmentArrives::Result::OK));
      test.execute(SegmentArrives{}.with_seqno(1).with_data("abcd").with_result(SegmentArrives::Result::OK));
      test.execute(ExpectAckno{WrappingInt32{5}});
      test.execute(ExpectBytes{"abcd"});
      test.execute(ExpectUnassembledBytes{0});
      test.execute(ExpectTotalAssembledBytes{4});
    }

    {
      uint32_t isn = 384678;
      TCPReceiverTestHarness test{4000};
      test.execute(SegmentArrives{}.with_syn().with_seqno(isn).with_result(SegmentArrives::Result::OK));
      test.execute(SegmentArrives{}.with_seqno(isn + 1).with_data("abcd").with_result(SegmentArrives::Result::OK));
      test.execute(ExpectAckno{WrappingInt32{isn + 5}});
      test.execute(ExpectUnassembledBytes{0});
      test.execute(ExpectTotalAssembledBytes{4});
      test.execute(ExpectBytes{"abcd"});
      test.execute(SegmentArrives{}.with_seqno(isn + 5).with_data("efgh").with_result(SegmentArrives::Result::OK));
      test.execute(ExpectAckno{WrappingInt32{isn + 9}});
      test.execute(ExpectUnassembledBytes{0});
      test.execute(ExpectTotalAssembledBytes{8});
      test.execute(ExpectBytes{"efgh"});
    }

    {
      uint32_t isn = 5;
      TCPReceiverTestHarness test{4000};
      test.execute(SegmentArrives{}.with_syn().with_seqno(isn).with_result(SegmentArrives::Result::OK));
      test.execute(SegmentArrives{}.with_seqno(isn + 1).with_data("abcd").with_result(SegmentArrives::Result::OK));
      test.execute(ExpectAckno{WrappingInt32{isn + 5}});
      test.execute(ExpectUnassembledBytes{0});
      test.execute(ExpectTotalAssembledBytes{4});
      test.execute(SegmentArrives{}.with_seqno(isn + 5).with_data("efgh").with_result(SegmentArrives::Result::OK));
      test.execute(ExpectAckno{WrappingInt32{isn + 9}});
      test.execute(ExpectUnassembledBytes{0});
      test.execute(ExpectTotalAssembledBytes{8});
      test.execute(ExpectBytes{"abcdefgh"});
    }

    {
      TCPReceiverTestHarness test{4000};
      test.execute(SegmentArrives{}.with_syn().with_seqno(0).with_result(SegmentArrives::Result::OK));
      test.execute(SegmentArrives{}.with_seqno(1).with_data("abcd").with_result(SegmentArrives::Result::OK));
      test.execute(SegmentArrives{}.with_seqno(1).with_data("efgh").with_result(SegmentArrives::Result::OUT_OF_WINDOW));
      test.execute(
          SegmentArrives{}.with_seqno(4005).with_data("efgh").with_result(SegmentArrives::Result::OUT_OF_WINDOW));
    }

    // Many (arrive/read)s
    {
      TCPReceiverTestHarness test{4000};
      uint32_t max_block_size = 10;
      uint32_t n_rounds = 10000;
      uint32_t isn = 893472;
      size_t bytes_sent = 0;
      test.execute(SegmentArrives{}.with_syn().with_seqno(isn).with_result(SegmentArrives::Result::OK));
      for (uint32_t i = 0; i < n_rounds; ++i) {
        string data;
        uint32_t block_size = uniform_int_distribution<uint32_t>{1, max_block_size}(rd);
        for (uint8_t j = 0; j < block_size; ++j) {
          uint8_t c = 'a' + ((i + j) % 26);
          data.push_back(c);
        }
        test.execute(ExpectAckno{WrappingInt32{isn + uint32_t(bytes_sent) + 1}});
        test.execute(ExpectTotalAssembledBytes{bytes_sent});
        test.execute(
            SegmentArrives{}.with_seqno(isn + bytes_sent + 1).with_data(data).with_result(SegmentArrives::Result::OK));
        bytes_sent += block_size;
        test.execute(ExpectBytes{std::move(data)});
      }
    }

    // Many arrives, one read
    {
      uint32_t max_block_size = 10;
      uint32_t n_rounds = 100;
      TCPReceiverTestHarness test{uint16_t(max_block_size * n_rounds)};
      uint32_t isn = 238;
      size_t bytes_sent = 0;
      test.execute(SegmentArrives{}.with_syn().with_seqno(isn).with_result(SegmentArrives::Result::OK));
      string all_data;
      for (uint32_t i = 0; i < n_rounds; ++i) {
        string data;
        uint32_t block_size = uniform_int_distribution<uint32_t>{1, max_block_size}(rd);
        for (uint8_t j = 0; j < block_size; ++j) {
          uint8_t c = 'a' + ((i + j) % 26);
          data.push_back(c);
          all_data.push_back(c);
        }
        test.execute(ExpectAckno{WrappingInt32{isn + uint32_t(bytes_sent) + 1}});
        test.execute(ExpectTotalAssembledBytes{bytes_sent});
        test.execute(
            SegmentArrives{}.with_seqno(isn + bytes_sent + 1).with_data(data).with_result(SegmentArrives::Result::OK));
        bytes_sent += block_size;
      }
      test.execute(ExpectBytes{std::move(all_data)});
    }

  } catch (const exception &e) {
    cerr << e.what() << endl;
    return 1;
  }

  return EXIT_SUCCESS;
}
